#include <iostream>
#include <vector>
#include <thread>
#include <windows.h>


typedef struct{
    int start;
    int end;
    int k;
    int n;
    int p2;
    int offset;
    float *Mix;
}Pot;

// 向量数乘
void vec_mul_s(float *A, float *B, float f, int size){
    int p = size;
    float *all = new float[size * 2];
    for (int i = 0; i < size; i++){
        all[i] = A[i];
    }

    for (int i = 0; i < size; i++){
        all[p + i] = f * all[i];
    }

    for (int i = 0; i < size; i++){
        B[i] = all[p + i];
    }
    delete[] all;
}

void vec_mul_t(Pot *p){
    for (int i = p->start; i < p->end; i++){
        p->Mix[i + p->offset] = p->Mix[i] * p->Mix[p->p2];
    }
}

void vec_mul_p(float *A, float *B, float f, int size){
    int p = size;
    int p2 = size * 2;
    float *all = new float[size * 2 + 1];
    for (int i = 0; i < size; i++){
        all[i] = A[i];
    }

    all[size * 2] = f;

    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot> mes(thread_num);

    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot{start, end, (int)NULL, (int)NULL, p2, size, all};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot{start, end, (int)NULL, (int)NULL, p2, size, all};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot{start, end, (int)NULL, (int)NULL, p2, size, all};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(vec_mul_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
    
    for (int i = 0; i < size; i++){
        B[i] = all[p + i];
    }
    delete[] all;
}

// 向量内积
void vec_dot_s(float *A, float *B, float *C, int size){
    int p1 = size;
    int p2 = p1 + size;
    float *all = new float[size * 3];
    for (int i = 0; i < size; i++){
        all[i] = A[i];
    }

    for (int i = 0; i < size; i++){
        all[p1 + i] = B[i];
    }

    for (int i = 0; i < size; i++){
        all[p2 + i] = all[i] * all[p1 + i];
    }

    for (int i = 0; i < size; i++){
        C[i] = all[p2 + i];
    }
    
    delete[] all;
}

void dsot_t(Pot *p){
    for (int i = p->start; i < p->end; i++){
        p->Mix[i + p->n] = p->Mix[i] * p->Mix[p->k + i];
    }
}

void vec_dot_p(float *A, float *B, float *C, int size){
    int p1 = size;
    int p2 = p1 + size;
    float *all = new float[size * 3];
    for (int i = 0; i < size; i++){
        all[i] = A[i];
    }

    for (int i = 0; i < size; i++){
        all[p1 + i] = B[i];
    }
    
    int thread_num = std::thread::hardware_concurrency();
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot> mes(thread_num);

    int blocksdim;
    if (size%thread_num != 0) {
        if (size < thread_num) {
            for (int i = 0; i < size; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot{start, end, p1, p2, (int)NULL, (int)NULL, all};
            }
        }else{
            blocksdim = size/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = size;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot{start, end, p1, p2, (int)NULL, (int)NULL, all};
            }
        }
    }else{
        blocksdim = size/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot{start, end, p1, p2, (int)NULL, (int)NULL, all};
            }
    }
    if (size < thread_num) {
        thread_num = size;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(dsot_t, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }
    
    for (int i = 0; i < size; i++){
        C[i] = all[p2 + i];
    }
    
    delete[] all;
}



// 非并行gemm
void gemm_s(float *A, float *B, float *C, int m, int n, int k){
    int p1 = m*k;
    int p2 = p1 + n*k;
    int all_size = p2 + m*n;

    float *all = new float[all_size];
    for (int i = 0; i < p1; i++){
        all[i] = A[i];
    }

    int idx = p1;
    for (int i = 0; i < n; i++){
        for (int j = 0; j < k; j++){
            all[idx] = B[j*n + i];
            idx++;
        }
    }
    
    idx = p2;
    for (int i = 0; i < m; i++){
        for (int j = m; j < n + m; j++){
            int o1 = i*k;
            int o2 = j*k;
            for (int c = 0; c < k; c++){
                all[idx] += all[o1+c] * all[o2+c];
            }
            idx++;
        }
    }
    for (int i = 0; i < m * n; i++){
        C[i] = all[p2 + i];
    }

    delete[] all;
}

// 并行式gemm

void gent(Pot *p){
    for (int f = p->start; f < p->end; f++){
        int i = f / p->n;
        int j = f % p->n;
        int o1 = i*p->k;
        int o2 = (j+p->offset)*p->k;
        int o3 = p->p2+f;
        for (int c = 0; c < p->k; c++){
            p->Mix[o3] += p->Mix[o1 + c] * p->Mix[o2 + c];
        }
    }
}

void gemm_p(float *A, float *B, float *C, int m, int n, int k){
    int nums = m*n;

    int p1 = m*k;
    int p2 = p1 + n*k;
    int all_size = p2 + m*n;

    float *all = new float[all_size];
    for (int i = 0; i < p1; i++){
        all[i] = A[i];
    }

    int idx = p1;
    for (int i = 0; i < n; i++){
        for (int j = 0; j < k; j++){
            all[idx] = B[j*n + i];
            idx++;
        }
    }

    // int thread_num = std::thread::hardware_concurrency();
    int thread_num = 16;
    std::vector<std::thread> threads(thread_num);
    std::vector<Pot> mes(thread_num);

    int blocksdim;
    if (nums%thread_num != 0) {
        if (nums < thread_num) {
            for (int i = 0; i < nums; i++) {
                int end = i + 1;
                int start = i;
                mes[i] = Pot{start, end, k, n, p2, m, all};
            }
        }else{
            blocksdim = nums/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = nums;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                mes[i] = Pot{start, end, k, n, p2, m, all};
            }
        }
    }else{
        blocksdim = nums/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                mes[i] = Pot{start, end, k, n, p2, m, all};
            }
    }
    if (nums < thread_num) {
        thread_num = nums;
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(gent, &mes[i]);
    }
    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }

    for (int i = 0; i < m * n; i++){
        C[i] = all[p2 + i];
    }

    delete[] all;
}

void gemm(float *A, float *B, float *C, int m, int n, int k){
    for (int i = 0; i < m; i++){
        for (int j = 0; j < n; j++){
            for (int c = 0; c < k; c++){
                C[i*n + j] += A[i*k + c] * B[c*n + j];
            }
        }
    }
}

void klm(){
    register int i, j, c;
}

int main(){
    int m = 512, n = 512, k = 512;
    float *A = new float[m*k];
    float *B = new float[k*n];
    float *C = new float[m*n];

    for (int i = 0; i < m*k; i++){
        A[i] = 1.0f;
    }

    for (int i = 0; i < k*n; i++){
        B[i] = 1.0f;
    }

    for (int i = 0; i < m*n; i++){
        C[i] = 0.0f;
    }

    DWORD start_t, end_t;

    start_t = GetTickCount();
    
    gemm_p(A, B, C, m , n, k);

    end_t = GetTickCount();

    double time = (end_t - start_t)/1e+3;

    std::cout << "time: " << time << "s" << std::endl;

    double flops = 2.0 * m * n * k * 1e-9;
    // std::cout << "flops: " << flops << "flops" << std::endl;
    std::cout << "flops: " << flops/time << "Gflops" << std::endl;


    // for (int i = 0; i < m; i++){
    //     for (int j = 0; j < n; j++){
    //         std::cout << C[i*n + j] << " ";
    //     }
    //     std::cout << std::endl;
    // }
    // for (int i = 0; i < 6; i++){
    //     std::cout << C[i] << " ";
    // }
    // std::cout << std::endl;

    delete[] A;
    delete[] B;
    delete[] C;
    return 0;
}